"""
限流中间件
提供基于令牌桶算法的请求限流功能
"""
import asyncio
import time
from typing import Dict, Optional
from collections import defaultdict, deque
from loguru import logger

from utils.logger import Logger


class RateLimiter:
    """限流器"""
    
    def __init__(self, redis_client):
        self.redis_client = redis_client
        self.logger = Logger.get_logger(agent_name="rate_limiter")
        
        # 本地限流缓存（用于快速检查）
        self.local_cache: Dict[str, deque] = defaultdict(deque)
        self.cache_size = 1000
        
        # 配置
        self.sync_interval = 1  # 秒，与Redis同步的间隔
        self.last_sync_time = time.time()
        
        self.logger.info("限流器初始化完成")
    
    async def is_allowed(self, client_id: str, limit: int, 
                        window_seconds: int = 60) -> bool:
        """
        检查是否允许请求
        
        Args:
            client_id: 客户端标识（用户ID或IP地址）
            limit: 限制数量
            window_seconds: 时间窗口（秒）
            
        Returns:
            bool: 是否允许请求
        """
        try:
            # 使用滑动窗口算法
            return await self._sliding_window_check(client_id, limit, window_seconds)
            
        except Exception as e:
            self.logger.error(f"限流检查失败: {str(e)}")
            # 失败时允许请求通过（fail-open）
            return True
    
    async def _sliding_window_check(self, client_id: str, limit: int, 
                                   window_seconds: int) -> bool:
        """滑动窗口限流检查"""
        try:
            current_time = time.time()
            
            # 获取请求历史
            request_history = await self._get_request_history(client_id, window_seconds)
            
            # 清理过期的请求记录
            cutoff_time = current_time - window_seconds
            valid_requests = [
                req_time for req_time in request_history 
                if req_time > cutoff_time
            ]
            
            # 检查是否超过限制
            if len(valid_requests) >= limit:
                self.logger.warning(
                    f"限流拒绝: client={client_id}, "
                    f"requests={len(valid_requests)}, limit={limit}"
                )
                return False
            
            # 记录当前请求
            valid_requests.append(current_time)
            await self._record_request(client_id, valid_requests, window_seconds)
            
            self.logger.debug(
                f"限流通过: client={client_id}, "
                f"requests={len(valid_requests)}, limit={limit}"
            )
            return True
            
        except Exception as e:
            self.logger.error(f"滑动窗口检查失败: {str(e)}")
            return True
    
    async def _get_request_history(self, client_id: str, 
                                  window_seconds: int) -> list:
        """获取请求历史"""
        try:
            # 先从本地缓存获取
            if client_id in self.local_cache:
                return list(self.local_cache[client_id])
            
            # 从Redis获取
            redis_key = f"rate_limit:requests:{client_id}"
            history_data = self.redis_client.get(redis_key)
            
            if history_data:
                import json
                history = json.loads(history_data)
                # 更新本地缓存
                self.local_cache[client_id] = deque(history, maxlen=self.cache_size)
                return history
            
            return []
            
        except Exception as e:
            self.logger.error(f"获取请求历史失败: {str(e)}")
            return []
    
    async def _record_request(self, client_id: str, request_history: list, 
                             window_seconds: int):
        """记录请求"""
        try:
            # 更新本地缓存
            self.local_cache[client_id] = deque(request_history, maxlen=self.cache_size)
            
            # 定期同步到Redis
            current_time = time.time()
            if current_time - self.last_sync_time >= self.sync_interval:
                await self._sync_to_redis()
                self.last_sync_time = current_time
            
        except Exception as e:
            self.logger.error(f"记录请求失败: {str(e)}")
    
    async def _sync_to_redis(self):
        """同步到Redis"""
        try:
            for client_id, history in self.local_cache.items():
                if not history:
                    continue
                
                redis_key = f"rate_limit:requests:{client_id}"
                
                # 序列化历史记录
                import json
                history_data = json.dumps(list(history))
                
                # 保存到Redis，设置过期时间
                ttl = 3600  # 1小时过期
                self.redis_client.setex(redis_key, ttl, history_data)
            
            self.logger.debug(f"限流数据同步完成，客户端数量: {len(self.local_cache)}")
            
        except Exception as e:
            self.logger.error(f"同步到Redis失败: {str(e)}")
    
    def get_rate_limit_stats(self, client_id: str, window_seconds: int = 60) -> dict:
        """获取限流统计信息"""
        try:
            # 获取请求历史
            request_history = self.local_cache.get(client_id, deque())
            
            current_time = time.time()
            cutoff_time = current_time - window_seconds
            
            # 计算有效请求
            valid_requests = [
                req_time for req_time in request_history 
                if req_time > cutoff_time
            ]
            
            return {
                "client_id": client_id,
                "window_seconds": window_seconds,
                "request_count": len(valid_requests),
                "first_request_time": min(valid_requests) if valid_requests else None,
                "last_request_time": max(valid_requests) if valid_requests else None,
                "cache_size": len(request_history)
            }
            
        except Exception as e:
            self.logger.error(f"获取限流统计失败: {str(e)}")
            return {
                "client_id": client_id,
                "error": str(e)
            }
    
    def clear_client_history(self, client_id: str):
        """清理客户端历史记录"""
        try:
            if client_id in self.local_cache:
                del self.local_cache[client_id]
            
            # 清理Redis中的记录
            redis_key = f"rate_limit:requests:{client_id}"
            self.redis_client.delete(redis_key)
            
            self.logger.info(f"客户端历史记录已清理: {client_id}")
            
        except Exception as e:
            self.logger.error(f"清理客户端历史记录失败: {str(e)}")


# --------------------------------------------------
# FastAPI 中间件封装
# --------------------------------------------------
from fastapi import Request, Response
from fastapi.responses import JSONResponse


async def rate_limiter_middleware(request: Request, call_next):
    """
    限流中间件
    依赖 Redis，需保证 request.app.state.redis 已挂载
    """
    redis = getattr(request.app.state, "redis", None)
    if not redis:
        logger.warning("[RateLimit] Redis 未挂载，跳过限流")
        return await call_next(request)

    # 按 IP 限流，默认 60 秒 120 次
    client_ip = request.client.host
    limiter = RateLimiter(redis)
    allowed = await limiter.is_allowed(client_ip, limit=120, window_seconds=60)

    if not allowed:
        return JSONResponse(
            status_code=429,
            content={"detail": "Too Many Requests"}
        )

    return await call_next(request)


class TokenBucketRateLimiter:
    """令牌桶限流器（备用实现）"""
    
    def __init__(self, redis_client):
        self.redis_client = redis_client
        self.logger = Logger.get_logger(agent_name="token_bucket_limiter")
    
    async def is_allowed(self, client_id: str, capacity: int, 
                        refill_rate: float) -> bool:
        """
        令牌桶限流检查
        
        Args:
            client_id: 客户端标识
            capacity: 桶容量
            refill_rate: 填充速率（令牌/秒）
            
        Returns:
            bool: 是否允许请求
        """
        try:
            key = f"token_bucket:{client_id}"
            current_time = time.time()
            
            # 获取当前状态
            bucket_data = self.redis_client.hmget(key, "tokens", "last_refill")
            tokens = float(bucket_data[0]) if bucket_data[0] else capacity
            last_refill = float(bucket_data[1]) if bucket_data[1] else current_time
            
            # 计算应该添加的令牌
            time_passed = current_time - last_refill
            tokens_to_add = time_passed * refill_rate
            tokens = min(capacity, tokens + tokens_to_add)
            
            # 检查是否有足够的令牌
            if tokens >= 1:
                # 消耗一个令牌
                tokens -= 1
                
                # 更新桶状态
                self.redis_client.hmset(key, {
                    "tokens": tokens,
                    "last_refill": current_time
                })
                
                # 设置过期时间（防止僵尸键）
                self.redis_client.expire(key, 3600)
                
                self.logger.debug(f"令牌桶通过: client={client_id}, tokens={tokens}")
                return True
            else:
                self.logger.warning(f"令牌桶拒绝: client={client_id}, tokens={tokens}")
                return False
                
        except Exception as e:
            self.logger.error(f"令牌桶检查失败: {str(e)}")
            return True